"""Tianyan platform utilities for quantum circuit generation and execution.

This module provides functions for generating benchmark circuits and handling
quantum circuit execution on the Tianyan platform.

Available quantum gates:
    - Single-qubit gates: 'rx', 'ry', 'rz', 'x', 'y', 'z', 'h'
    - Two-qubit gates: 'cx', 'cz'
    - Other operations: 'delay', 'barrier', 'measure'
"""

import json
from datetime import datetime
from typing import List, Dict, Tuple, Optional, Any, Union

from cqlib import TianYanPlatform, QuantumLanguage
from cqlib.circuits import Circuit
from cqlib.utils import LaboratoryUtils


def send_job(circuit: Circuit,
            lab_id: str,
            nshots: int,
            platform: TianYanPlatform) -> str:
    """Submit job to Tianyan platform.

    Args:
        circuit: Quantum circuit to execute
        lab_id: Laboratory ID
        nshots: Number of shots
        platform: TianYan platform instance

    Returns:
        Query ID for the job
    """
    return platform.submit_job(
        circuit=circuit.qcis,
        exp_name=f'exp.{datetime.now().strftime("%Y%m%d%H%M%S")}',
        lab_id=lab_id,
        num_shots=nshots
    )


def get_result_less_15qubits(df: Any, 
                            platform: TianYanPlatform) -> None:
    """Get results for circuits with less than 15 qubits.

    Args:
        df: Dataframe containing experiment data
        platform: TianYan platform instance
    """
    lab_utils = LaboratoryUtils()
    
    if 'calibrated_exp_result' not in df.columns:
        df['calibrated_exp_result'] = None

    for row in range(len(df)):
        exp_id = df.loc[row, 'exp_id']
        if '[' in exp_id:
            exp_id = eval(exp_id)[0]
            
        result = platform.query_experiment(
            query_id=exp_id,
            max_wait_time=120,
            sleep_time=5
        )
        
        df.loc[row, 'exp_result'] = json.dumps(
            json.loads(result[0]['probability'])
        )

        # Apply calibration and correction
        calibration_result = lab_utils.probability_calibration(
            result=result[0],
            laboratory=platform
        )
        corrected_result = lab_utils.probability_correction(calibration_result)
        df.loc[row, 'calibrated_exp_result'] = json.dumps(corrected_result)


def get_result_more_15qubits(df: Any, 
                            platform: TianYanPlatform) -> None:
    """Get results for circuits with more than 15 qubits.

    Args:
        df: Dataframe containing experiment data
        platform: TianYan platform instance
    """
    lab_utils = LaboratoryUtils()
    
    if 'calibrated_exp_result' not in df.columns:
        df['calibrated_exp_result'] = None

    for row in range(len(df)):
        exp_id = df.loc[row, 'exp_id']
        if '[' in exp_id:
            exp_id = eval(exp_id)[0]
            
        result = platform.query_experiment(
            query_id=exp_id,
            max_wait_time=120,
            sleep_time=5
        )
        
        probability = lab_utils.readout_data_to_state_probabilities_part(result[0])
        df.loc[row, 'exp_result'] = json.dumps(probability)

        # Apply calibration and correction
        calibration_result = lab_utils.probability_calibration(
            result=result[0],
            laboratory=platform
        )
        corrected_result = lab_utils.probability_correction(calibration_result)
        df.loc[row, 'calibrated_exp_result'] = json.dumps(corrected_result)
